import sys
from time import time
sys.path.insert(0,'..')
from rmpoly import *
from rmpoly.free_algebra import *
from gmpy import mpq

# example from talk by F.Casas "On the BCH formula and related expansions"
def bch1(n,h):
  """
  log(exp(x[0]*A)*exp(x[1]*B)*...*exp(x[2*n-2]*A)*exp(x[2*n-1]*B))
  """
  fa,A,B = fgens('A,B',mpq)
  rp = NCRPoly(['x%d' %i for i in range(2*n)],8,fa,order='grlex')
  x = rp.gens()
  t0 = time()
  p = ((A*x[0]).exp('_t',h)).mul_trunc((B*x[1]).exp('_t',h),'_t',h)
  for i in range(2,n,2):
    p1 = ((A*x[i]).exp('_t',h)).mul_trunc((B*x[i+1]).exp('_t',h),'_t',h)
    p = p.mul_trunc(p1,'_t',h)
  p = p.log('_t',h)
  t1 = time()
  print 'p=',p
  print '%.2f' %(t1-t0)


# http://mathworld.wolfram.com/Baker-Campbell-HausdorffSeries.html
def bch2(n,h):
  """
  log(exp(x[0]*A)*...*exp(x[n-1]*A))
  """
  fa = FreeAlgebra(['A%d' %i for i in range(n)], mpq)
  A = fa.gens()
  rp,x = ncrgens('x',8,fa)
  p = (A[0]*x).exp('x',h)
  for i in range(1,n):
    p = p.mul_trunc( (A[i]*x).exp('x',h), 'x',h)
  p = p.log('x',h)
  print p



def state_evolution(n,ht):
  """
  probability amplitude for transition from the vacuum state
  to the state with occupation number n
  for the anharmonic oscillator
  <n| exp(t*H)|0>
  """
  fa,A,Ad,Bra,Ket = fgens('A,Ad,Bra,Ket',mpq)
  # FIXME writing Bra*Ket, 1 it fails with exception
  fa.add_slice_rules(A*Ket,0, Bra*Ad,0, A*Ad,Ad*A+1, Bra*Ket,fa(1))
  fa.set_rules()
  rp,t,g = ncrgens('t,g',8,fa)
  #H = Ad*A + g*(Ad + A)**4
  #print 'H=',H
  # normal ordered H
  H=  Ad*A + g*(Ad**4 +4*Ad**3*A +6*Ad**2*A**2 +4*Ad*A**3 +A**4)
  p = Bra * A**n * (t*H).exp('t',ht) * Ket
  print 'p=',p


#bch1(1,10)
#bch2(3,5)
state_evolution(4,6)


